Add D256 to GQA-8 two-pass vector attention - #4477

import argparse
import statistics
import sys
import time
from pathlib import Path

p = argparse.ArgumentParser()
p.add_argument('--package', required=True, help='Source build python/ directory')
p.add_argument('--dtype', default='float16')
p.add_argument('--q', type=int, required=True)
p.add_argument('--k', type=int, required=True)
p.add_argument('--batch', type=int, default=1)
p.add_argument('--hq', type=int, default=16)
p.add_argument('--hk', type=int, default=2)
p.add_argument('--causal', action='store_true')
a = p.parse_args()
a.package = str(Path(a.package).resolve())
sys.path.insert(0, a.package)
import mlx.core as mx

assert a.package in mx.__file__, mx.__file__
print("BUILD", mx.__file__)

mx.set_default_device(mx.gpu)
mx.random.seed(0)
dtype = getattr(mx, a.dtype)
q = mx.random.normal((a.batch, a.hq, a.q, 256)).astype(dtype)
k = mx.random.normal((a.batch, a.hk, a.k + 32, 256)).astype(dtype)[:, :, :a.k]
v = mx.random.normal((a.batch, a.hk, a.k + 32, 256)).astype(dtype)[:, :, :a.k]
mx.eval(q, k, v)

def chain():
    x = q
    for _ in range(4):
        x = mx.fast.scaled_dot_product_attention(
            x, k, v, scale=1 / 16, mask='causal' if a.causal else None
        )
    mx.eval(x)

for _ in range(8):
    chain()
mx.synchronize()
samples = []
for _ in range(5):
    start = time.perf_counter()
    for _ in range(64):
        chain()
    mx.synchronize()
    samples.append((time.perf_counter() - start) / (64 * 4))
print('RESULT', statistics.median(samples))